# -*-coding: utf-8 -*-
"""
    @Theme   : MongoDB Client API
    @Time    : 2024/7/7 14:08
    @Author  : YamFish
    @Url     : https://www.zhihu.com/people/yamfish
"""
from typing import Dict, Optional

import pymongo
from bson import ObjectId
from pymongo.cursor import Cursor
from pymongo.results import UpdateResult
from pymongo.collection import Collection
from pymongo.typings import _DocumentType
from pymongo.command_cursor import CommandCursor


class MongoClient:
    def __init__(self, app=None, uri=None, db=None):
        if app:
            self.db = pymongo.MongoClient(app.config['MONGODB_URI'])[app.config['MONGODB_DATABASE']]
        else:
            self.db = pymongo.MongoClient(uri)[db]
        self.handler_dict: Dict[str, Collection] = {}

    def set_handler(self, collection: str):
        self.handler_dict[collection] = self.db[collection]

    def insert_one(self, collection: str, data: dict) -> ObjectId:
        object_id = self.handler_dict[collection].insert_one(data).inserted_id
        return object_id

    def find_one(self, collection: str, *args) -> Optional[_DocumentType]:
        return self.handler_dict[collection].find_one(*args)

    def find(self, collection: str, *args) -> Cursor:
        cursor = self.handler_dict[collection].find(*args)
        return cursor

    def update_one(self, collection: str, *args) -> UpdateResult:
        return self.handler_dict[collection].update_one(*args)

    def aggregate(self, collection: str, *args) -> CommandCursor:
        return self.handler_dict[collection].aggregate(*args)

    def get_count(self, collection: str, *args) -> int:
        return self.handler_dict[collection].count_documents(*args)
